package org.github.manx98.jlibhac.tools.fssystem;

import cn.hutool.core.util.ArrayUtil;
import cn.hutool.crypto.SecureUtil;
import org.github.manx98.jlibhac.common.Utilities;
import org.github.manx98.jlibhac.common.Validity;
import org.github.manx98.jlibhac.fs.IStorage;

import java.io.IOException;

/**
 * @author Manx98
 * @version 1.0
 * @since 2022/10/4 0:22
 */
public class IntegrityVerificationStorage extends SectorStorage {
    private static final int DigestSize = 0x20;
    private final IStorage hashStorage;
    private final IntegrityCheckLevel integrityCheckLevel;
    private final byte[] salt;
    private final IntegrityStorageType type;
    private final Validity[] blockValidities;

    public IntegrityVerificationStorage(IntegrityVerificationInfo info, IStorage hashStorage, IntegrityCheckLevel integrityCheckLevel, boolean leaveOpen) {
        super(info.data, info.blockSize, leaveOpen);
        this.hashStorage = hashStorage;
        this.integrityCheckLevel = integrityCheckLevel;
        this.salt = info.salt;
        this.type = info.type;
        this.blockValidities = new Validity[sectorCount];
    }

    public Validity[] getBlockValidities() {
        return blockValidities;
    }

    private void readImpl(long offset, byte[] destination, int destOffset, int destLen, IntegrityCheckLevel integrityCheckLevel) throws IOException {

        if (destLen < 0 || destLen > sectorSize)
            throw new IllegalArgumentException("Length " + destLen + " is invalid, must be between 0 and " + sectorSize);

        int blockIndex = Math.toIntExact(offset / sectorSize);

        if (blockValidities[blockIndex] == Validity.Invalid && integrityCheckLevel == IntegrityCheckLevel.ErrorOnInvalid) {
            // Todo: Differentiate between the top and lower layers
            throw new IllegalArgumentException("Hash error!");
        }

        boolean needsHashCheck = integrityCheckLevel != IntegrityCheckLevel.None &&
                blockValidities[blockIndex] == Validity.Unchecked;

        if (type != IntegrityStorageType.Save && !needsHashCheck) {
            baseStorage.read(offset, destination, destOffset, destLen);
            return;
        }

        byte[] hashBuffer = new byte[DigestSize];
        long hashPos = (long) blockIndex * DigestSize;
        hashStorage.read(hashPos, hashBuffer, 0, hashBuffer.length);

        if (type == IntegrityStorageType.Save) {
            if (Utilities.isZeros(hashBuffer)) {
                blockValidities[blockIndex] = Validity.Valid;
                return;
            }
            if (!needsHashCheck) {
                baseStorage.read(offset, destination, destOffset, destLen);
                return;
            }
        }

        byte[] dataBuffer = new byte[sectorSize];
        baseStorage.read(offset, destination, destOffset, destLen);
        System.arraycopy(dataBuffer, 0, dataBuffer, 0, dataBuffer.length);

        if (blockValidities[blockIndex] != Validity.Unchecked) return;

        int bytesToHash = sectorSize;

        if (destLen < sectorSize) {
            // Pad out unused portion of block
            for (int i = destLen; i < sectorSize - destLen; ++i) {
                dataBuffer[i] = 0;
            }
            // Partition FS hashes don't pad out an incomplete block
            if (type == IntegrityStorageType.PartitionFs) {
                bytesToHash = destLen;
            }
        }

        byte[] hash = SecureUtil.sha256().digest(dataBuffer);

        Validity validity = Utilities.spansEqual(hashBuffer, 0, hash, 0, hash.length) ? Validity.Valid : Validity.Invalid;
        blockValidities[blockIndex] = validity;

        if (validity == Validity.Invalid && integrityCheckLevel == IntegrityCheckLevel.ErrorOnInvalid) {
            throw new IllegalStateException("Hash error!");
        }
    }

    @Override
    public void read(long offset, byte[] destination, int destOffset, int len) throws IOException {
        readImpl(offset, destination, destOffset, len, integrityCheckLevel);
    }
}
